package com.element.oauth2.server.global.provider;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.core.convert.TypeDescriptor;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2Token;
import org.springframework.security.oauth2.core.converter.ClaimConversionService;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenIntrospection;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2TokenIntrospectionAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import java.net.URL;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class OAuthTokenIntrospectionAuthenticationProvider implements AuthenticationProvider {

    private static final TypeDescriptor OBJECT_TYPE_DESCRIPTOR = TypeDescriptor.valueOf(Object.class);
    private static final TypeDescriptor LIST_STRING_TYPE_DESCRIPTOR = TypeDescriptor.collection(List.class, TypeDescriptor.valueOf(String.class));
    private final Log logger = LogFactory.getLog(this.getClass());
    private final RegisteredClientRepository registeredClientRepository;
    private final OAuth2AuthorizationService authorizationService;

    public OAuthTokenIntrospectionAuthenticationProvider(RegisteredClientRepository registeredClientRepository, OAuth2AuthorizationService authorizationService) {
        Assert.notNull(registeredClientRepository, "registeredClientRepository cannot be null");
        Assert.notNull(authorizationService, "authorizationService cannot be null");
        this.registeredClientRepository = registeredClientRepository;
        this.authorizationService = authorizationService;
    }

    public Authentication authenticate(Authentication authentication) throws AuthenticationException {
        OAuth2TokenIntrospectionAuthenticationToken tokenIntrospectionAuthentication = (OAuth2TokenIntrospectionAuthenticationToken) authentication;
        OAuth2ClientAuthenticationToken clientPrincipal = getAuthenticatedClientElseThrowInvalidClient(tokenIntrospectionAuthentication);
        OAuth2TokenType tokenType = null;
        if (OAuth2TokenType.ACCESS_TOKEN.getValue().equals(tokenIntrospectionAuthentication.getTokenTypeHint())) {
            tokenType = OAuth2TokenType.ACCESS_TOKEN;
        } else if (OAuth2TokenType.REFRESH_TOKEN.getValue().equals(tokenIntrospectionAuthentication.getTokenTypeHint())) {
            tokenType = OAuth2TokenType.REFRESH_TOKEN;
        }
        OAuth2Authorization authorization = this.authorizationService.findByToken(tokenIntrospectionAuthentication.getToken(), tokenType);
        if (authorization == null) {
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Did not authenticate token introspection request since token was not found");
            }

            return tokenIntrospectionAuthentication;
        } else {
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Retrieved authorization with token");
            }

            OAuth2Authorization.Token<OAuth2Token> authorizedToken = authorization.getToken(tokenIntrospectionAuthentication.getToken());
            assert authorizedToken != null;
            if (!authorizedToken.isActive()) {
                if (this.logger.isTraceEnabled()) {
                    this.logger.trace("Did not introspect token since not active");
                }

                return new OAuth2TokenIntrospectionAuthenticationToken(tokenIntrospectionAuthentication.getToken(), clientPrincipal, OAuth2TokenIntrospection.builder().build());
            } else {
                RegisteredClient authorizedClient = this.registeredClientRepository.findByClientId(authorization.getRegisteredClientId());
                OAuth2TokenIntrospection tokenClaims = withActiveTokenClaims(authorizedToken, authorizedClient);
                if (this.logger.isTraceEnabled()) {
                    this.logger.trace("Authenticated token introspection request");
                }

                return new OAuth2TokenIntrospectionAuthenticationToken(authorizedToken.getToken().getTokenValue(), clientPrincipal, tokenClaims);
            }
        }
    }

    public boolean supports(Class<?> authentication) {
        return OAuth2TokenIntrospectionAuthenticationToken.class.isAssignableFrom(authentication);
    }

    private OAuth2TokenIntrospection withActiveTokenClaims(OAuth2Authorization.Token<OAuth2Token> authorizedToken, RegisteredClient authorizedClient) {
        OAuth2TokenIntrospection.Builder tokenClaims;
        if (!CollectionUtils.isEmpty(authorizedToken.getClaims())) {
            Map<String, Object> claims = convertClaimsIfNecessary(authorizedToken.getClaims());
            tokenClaims = OAuth2TokenIntrospection.withClaims(claims).active(true);
        } else {
            tokenClaims = OAuth2TokenIntrospection.builder(true);
        }

        tokenClaims.clientId(authorizedClient.getClientId());
        OAuth2Token token = authorizedToken.getToken();
        if (token.getIssuedAt() != null) {
            tokenClaims.issuedAt(token.getIssuedAt());
        }

        if (token.getExpiresAt() != null) {
            tokenClaims.expiresAt(token.getExpiresAt());
        }

        if (OAuth2AccessToken.class.isAssignableFrom(token.getClass())) {
            OAuth2AccessToken accessToken = (OAuth2AccessToken) token;
            tokenClaims.tokenType(accessToken.getTokenType().getValue());
        }

        return tokenClaims.build();
    }

    private Map<String, Object> convertClaimsIfNecessary(Map<String, Object> claims) {
        Map<String, Object> convertedClaims = new HashMap<>(claims);
        Object value = claims.get("iss");
        if (value != null && !(value instanceof URL)) {
            URL convertedValue = ClaimConversionService.getSharedInstance().convert(value, URL.class);
            if (convertedValue != null) {
                convertedClaims.put("iss", convertedValue);
            }
        }

        value = claims.get("scope");
        Object convertedValue;
        if (value != null && !(value instanceof List)) {
            convertedValue = ClaimConversionService.getSharedInstance().convert(value, OBJECT_TYPE_DESCRIPTOR, LIST_STRING_TYPE_DESCRIPTOR);
            if (convertedValue != null) {
                convertedClaims.put("scope", convertedValue);
            }
        }

        value = claims.get("aud");
        if (value != null && !(value instanceof List)) {
            convertedValue = ClaimConversionService.getSharedInstance().convert(value, OBJECT_TYPE_DESCRIPTOR, LIST_STRING_TYPE_DESCRIPTOR);
            if (convertedValue != null) {
                convertedClaims.put("aud", convertedValue);
            }
        }

        return convertedClaims;
    }

    private OAuth2ClientAuthenticationToken getAuthenticatedClientElseThrowInvalidClient(Authentication authentication) {
        OAuth2ClientAuthenticationToken clientPrincipal = null;
        if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(authentication.getPrincipal().getClass())) {
            clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
        }

        if (clientPrincipal != null && clientPrincipal.isAuthenticated()) {
            return clientPrincipal;
        } else {
            throw new OAuth2AuthenticationException("invalid_client");
        }
    }
}